from torch.utils.data import Dataset
import json
import os
from pathlib import Path


class AFQMC(Dataset):
    # 初始化数据集参数
    def __init__(self, data_file):
        # 检查文件是否存在
        if not os.path.exists(data_file):
            raise FileNotFoundError(f"数据文件不存在: {data_file}")
        self.data = self.load_data(data_file)

    def load_data(self, data_file):
        Data = {}
        try:
            with open(data_file, 'rt', encoding='utf-8') as f:
                for idx, line in enumerate(f):
                    sample = json.loads(line.strip())
                    Data[idx] = sample
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
            print(f"错误位置: 文件 {data_file} 的第 {idx + 1} 行")
            raise
        return Data

    # 数据集中样本的个数
    def __len__(self):
        return len(self.data)

    # 映射型数据集的核心，根据给定的索引 idx 返回样本。
    def __getitem__(self, idx):
        return self.data[idx]


# 检查并创建数据目录
data_dir = Path('./data/afqmc_public')
data_dir.mkdir(parents=True, exist_ok=True)

# 示例: 加载数据
try:
    # 注意: 请确保文件实际存在于指定路径
    train_data = AFQMC('./data/afqmc_public/train.json')
    valid_data = AFQMC('./data/afqmc_public/dev.json')

    print(f"训练集大小: {len(train_data)}")
    print(f"验证集大小: {len(valid_data)}")
    print("\n第一个训练样本:")
    print(train_data[0])

except FileNotFoundError as e:
    print(f"文件不存在: {e}")
    print("请确认:")
    print("1. 数据文件是否已下载并放置在正确路径")
    print("2. 当前工作目录是否正确 (使用 os.getcwd() 查看)")
except Exception as e:
    print(f"其他错误: {e}")